import random, re, nltk

# ! Посмотреть про цепи Маркова и реализовать подобную тему. Может быть, это улучшит выходные результаты.

def tokenize(arr):
    """
    Токенизацция текста (получение массива всех слов) в массиве строк arr
    """
    return [w for l in arr for w in l.split(' ')]

def do_model(t):
    """
    Подготовка биграммной модели из массива строк текста t
    """
    # ! Использовать инструменты для естественного анализа? nltk и всё такое..
    t = tokenize(t)
    ngrams = {}

    for i in range(len(t)-2):
        seq = ' '.join(t[i:i+2])
        if seq not in ngrams.keys():
            ngrams[seq] = []
        ngrams[seq].append(t[i+2])
        
    return ngrams


def do_text(t: dict, SEQ):
    """
    Делает текст случайной длины с помощью биграммного словаря t с начальным запросом SEQ
    """
    output = SEQ
    for _ in range(random.randint(10, 24)):
        if SEQ not in t.keys():
            break
        possible_words = t[SEQ]
        next_word = possible_words[random.randrange(0, len(possible_words))]
        output += ' ' + next_word
        seq_words = nltk.word_tokenize(output, language="russian")
        SEQ = ' '.join(seq_words[len(seq_words)-2:len(seq_words)])
    
    return output

def prepare_text(text: list[str]):
    """
    Подготавливает текст к программной обработке
    * удаление двойных пробелов
    * удаление знаков препинания в тексте
    * перевод всего текста в нижний регистр

    Принимает и возвращает массив строк текста
    """
    return [re.sub(r'\s+', ' ', re.sub(r'[^\w\s]', '', x.lower())) for x in text]
